
#include "scatter_elements_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  ScatterElementsTilingData tiling;

    auto shape_var = context->GetInputShape(0)->GetOriginShape();
    auto shape_indices = context->GetInputShape(1)->GetOriginShape();
    auto shape_updates = context->GetInputShape(2)->GetOriginShape();

    // int32_t indices_ndarray[20];
    // int32_t indices_dimensional;

    // indices_dimensional = shape_indices.GetDimNum();
    // for(int i = 0; i < indices_dimensional; i++)
    // {
    //     indices_ndarray[i] = shape_indices.GetDim(i);
    // }

    int32_t reduce = 0;
    const char * reduce_str = context->GetAttrs()->GetAttrPointer<char>(1);;
    if(strcmp(reduce_str, "add") == 0)
    {
        reduce = 1;
    }
    else if(strcmp(reduce_str, "multiply") == 0)
    {
        reduce = 2;
    }
    else if(strcmp(reduce_str, "None") == 0)
    {
        reduce = 0;
    }
    else if(strcmp(reduce_str, "none") == 0)
    {
        reduce = 0;
    }
    else
    {
        reduce = 3;
    }
    
    int32_t axis = *context->GetAttrs()->GetInt(0);
    if(axis < 0) axis += shape_indices.GetDimNum();

    
    int32_t var_dim[3] = {1,1,1};
    int32_t indices_dim[3] = {1,1,1};
    int32_t updates_dim[3] = {1,1,1};

    for(int i = 0; i < axis; i++)
    {
        indices_dim[0] *= shape_indices.GetDim(i);
        var_dim[0] *= shape_var.GetDim(i);
        updates_dim[0] *= shape_updates.GetDim(i);
    }    
    indices_dim[1] = shape_indices.GetDim(axis);
    var_dim[1] = shape_var.GetDim(axis);
    updates_dim[1] = shape_updates.GetDim(axis);
    
    for(int i = axis+1; i < shape_indices.GetDimNum(); i++)
    {
        indices_dim[2] *= shape_indices.GetDim(i);
        var_dim[2] *= shape_var.GetDim(i);
        updates_dim[2] *= shape_updates.GetDim(i);
    }


    tiling.set_var_dim(var_dim);
    tiling.set_indices_dim(indices_dim);
    tiling.set_updates_dim(updates_dim);
    tiling.set_reduce(reduce);

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ScatterElements : public OpDef {
public:
    explicit ScatterElements(const char* name) : OpDef(name)
    {
        this->Input("var")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("updates")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("axis").AttrType(OPTIONAL).Int(0);
        this->Attr("reduce").AttrType(OPTIONAL).String("None");

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(ScatterElements);
}
